{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "k0_-IVVOk-9w",
    "outputId": "4a4c48bf-9257-406b-9a04-2071549095d4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
      "Requirement already satisfied: tensorflow in /usr/local/lib/python3.8/dist-packages (2.11.0)\n",
      "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (0.2.0)\n",
      "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (3.3.0)\n",
      "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (1.50.0)\n",
      "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (2.1.1)\n",
      "Requirement already satisfied: absl-py>=1.0.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (1.3.0)\n",
      "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (1.15.0)\n",
      "Requirement already satisfied: wrapt>=1.11.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (1.14.1)\n",
      "Requirement already satisfied: protobuf<3.20,>=3.9.2 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (3.19.6)\n",
      "Requirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from tensorflow) (57.4.0)\n",
      "Requirement already satisfied: tensorflow-estimator<2.12,>=2.11.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (2.11.0)\n",
      "Requirement already satisfied: tensorboard<2.12,>=2.11 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (2.11.0)\n",
      "Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (3.1.0)\n",
      "Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (4.1.1)\n",
      "Requirement already satisfied: numpy>=1.20 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (1.21.6)\n",
      "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (1.6.3)\n",
      "Requirement already satisfied: flatbuffers>=2.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (22.11.23)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from tensorflow) (21.3)\n",
      "Requirement already satisfied: libclang>=13.0.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (14.0.6)\n",
      "Requirement already satisfied: keras<2.12,>=2.11.0 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (2.11.0)\n",
      "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (0.28.0)\n",
      "Requirement already satisfied: gast<=0.4.0,>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from tensorflow) (0.4.0)\n",
      "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.8/dist-packages (from astunparse>=1.6.0->tensorflow) (0.38.4)\n",
      "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.12,>=2.11->tensorflow) (1.8.1)\n",
      "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.12,>=2.11->tensorflow) (2.14.1)\n",
      "Requirement already satisfied: requests<3,>=2.21.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.12,>=2.11->tensorflow) (2.23.0)\n",
      "Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.12,>=2.11->tensorflow) (1.0.1)\n",
      "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.12,>=2.11->tensorflow) (0.6.1)\n",
      "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.12,>=2.11->tensorflow) (3.4.1)\n",
      "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.8/dist-packages (from tensorboard<2.12,>=2.11->tensorflow) (0.4.6)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow) (0.2.8)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow) (4.9)\n",
      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow) (5.2.0)\n",
      "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.8/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.12,>=2.11->tensorflow) (1.3.1)\n",
      "Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.8/dist-packages (from markdown>=2.6.8->tensorboard<2.12,>=2.11->tensorflow) (4.13.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard<2.12,>=2.11->tensorflow) (3.10.0)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.8/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.12,>=2.11->tensorflow) (0.4.8)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard<2.12,>=2.11->tensorflow) (2022.9.24)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard<2.12,>=2.11->tensorflow) (2.10)\n",
      "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard<2.12,>=2.11->tensorflow) (1.24.3)\n",
      "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests<3,>=2.21.0->tensorboard<2.12,>=2.11->tensorflow) (3.0.4)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.8/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard<2.12,>=2.11->tensorflow) (3.2.2)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging->tensorflow) (3.0.9)\n",
      "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
      "Collecting git+https://github.com/DavidLandup0/keras-cv@vit\n",
      "  Cloning https://github.com/DavidLandup0/keras-cv (to revision vit) to /tmp/pip-req-build-76ui5l8a\n",
      "  Running command git clone -q https://github.com/DavidLandup0/keras-cv /tmp/pip-req-build-76ui5l8a\n",
      "  Running command git checkout -b vit --track origin/vit\n",
      "  Switched to a new branch 'vit'\n",
      "  Branch 'vit' set up to track remote branch 'vit' from 'origin'.\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.8/dist-packages (from keras-cv==0.3.4) (21.3)\n",
      "Requirement already satisfied: absl-py in /usr/local/lib/python3.8/dist-packages (from keras-cv==0.3.4) (1.3.0)\n",
      "Requirement already satisfied: regex in /usr/local/lib/python3.8/dist-packages (from keras-cv==0.3.4) (2022.6.2)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging->keras-cv==0.3.4) (3.0.9)\n"
     ]
    }
   ],
   "source": [
    "!pip install --upgrade tensorflow\n",
    "!pip install \"git+https://github.com/DavidLandup0/keras-cv@vit\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "m_v5YoiblKQw"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import sys\n",
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "import tensorflow as tf\n",
    "import keras_cv\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "E2aRsSNIl8I8"
   },
   "outputs": [],
   "source": [
    "config_names = {\n",
    "    # JAX name : KCV name, layer_nums\n",
    "    \"Ti/16\": (\"ViTTiny16\", 12),\n",
    "    \"S/16\": (\"ViTS16\", 12),\n",
    "    \"B/16\": (\"ViTB16\", 12),\n",
    "    \"L/16\": (\"ViTL16\", 24),\n",
    "    \"S/32\": (\"ViTS32\", 12),\n",
    "    \"B/32\": (\"ViTB32\", 12),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Xl-x6Uv4mqaI",
    "outputId": "84eb0986-9fe4-4ae1-e583-6d3caac8d501"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('Ti/16', ('ViT_Tiny_16', 12))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Choose model to convert\n",
    "model_to_convert = list(config_names.items())[0]\n",
    "model_to_convert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "P3Yf-sghlAU_"
   },
   "outputs": [],
   "source": [
    "model = eval(\n",
    "    f\"keras_cv.models.{model_to_convert[1][0]}(include_rescaling=False, include_top=True, num_classes=1000, weights=None, input_shape=(224, 224, 3))\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 733
    },
    "id": "N0riArBzlBF7",
    "outputId": "10151100-44c9-4413-c4ac-4dd5d2b764e6"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <div id=\"df-71a9d602-556d-4ab3-9618-c6628552dd1a\">\n",
       "    <div class=\"colab-df-container\">\n",
       "      <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>ds</th>\n",
       "      <th>epochs</th>\n",
       "      <th>lr</th>\n",
       "      <th>aug</th>\n",
       "      <th>wd</th>\n",
       "      <th>do</th>\n",
       "      <th>sd</th>\n",
       "      <th>best_val</th>\n",
       "      <th>final_val</th>\n",
       "      <th>...</th>\n",
       "      <th>adapt_ds</th>\n",
       "      <th>adapt_lr</th>\n",
       "      <th>adapt_steps</th>\n",
       "      <th>adapt_resolution</th>\n",
       "      <th>adapt_final_val</th>\n",
       "      <th>adapt_final_test</th>\n",
       "      <th>params</th>\n",
       "      <th>infer_samples_per_sec</th>\n",
       "      <th>filename</th>\n",
       "      <th>adapt_filename</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i1k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>light0</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.702544</td>\n",
       "      <td>0.702232</td>\n",
       "      <td>...</td>\n",
       "      <td>imagenet2012</td>\n",
       "      <td>0.030</td>\n",
       "      <td>20000</td>\n",
       "      <td>384</td>\n",
       "      <td>0.755698</td>\n",
       "      <td>0.72874</td>\n",
       "      <td>5790000.0</td>\n",
       "      <td>609.58</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i1k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>light0</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.702544</td>\n",
       "      <td>0.702232</td>\n",
       "      <td>...</td>\n",
       "      <td>imagenet2012</td>\n",
       "      <td>0.010</td>\n",
       "      <td>20000</td>\n",
       "      <td>384</td>\n",
       "      <td>0.754605</td>\n",
       "      <td>0.72412</td>\n",
       "      <td>5790000.0</td>\n",
       "      <td>609.58</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i1k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>light0</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.702544</td>\n",
       "      <td>0.702232</td>\n",
       "      <td>...</td>\n",
       "      <td>cifar100</td>\n",
       "      <td>0.030</td>\n",
       "      <td>10000</td>\n",
       "      <td>384</td>\n",
       "      <td>0.836000</td>\n",
       "      <td>0.83380</td>\n",
       "      <td>5790000.0</td>\n",
       "      <td>609.58</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i1k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>light0</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.702544</td>\n",
       "      <td>0.702232</td>\n",
       "      <td>...</td>\n",
       "      <td>cifar100</td>\n",
       "      <td>0.010</td>\n",
       "      <td>10000</td>\n",
       "      <td>384</td>\n",
       "      <td>0.835000</td>\n",
       "      <td>0.83040</td>\n",
       "      <td>5790000.0</td>\n",
       "      <td>609.58</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i1k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>light0</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.1</td>\n",
       "      <td>0.702544</td>\n",
       "      <td>0.702232</td>\n",
       "      <td>...</td>\n",
       "      <td>cifar100</td>\n",
       "      <td>0.003</td>\n",
       "      <td>10000</td>\n",
       "      <td>384</td>\n",
       "      <td>0.800000</td>\n",
       "      <td>0.79620</td>\n",
       "      <td>5790000.0</td>\n",
       "      <td>609.58</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "      <td>Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 21 columns</p>\n",
       "</div>\n",
       "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-71a9d602-556d-4ab3-9618-c6628552dd1a')\"\n",
       "              title=\"Convert this dataframe to an interactive table.\"\n",
       "              style=\"display:none;\">\n",
       "        \n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "       width=\"24px\">\n",
       "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
       "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
       "  </svg>\n",
       "      </button>\n",
       "      \n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      flex-wrap:wrap;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "      <script>\n",
       "        const buttonEl =\n",
       "          document.querySelector('#df-71a9d602-556d-4ab3-9618-c6628552dd1a button.colab-df-convert');\n",
       "        buttonEl.style.display =\n",
       "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "        async function convertToInteractive(key) {\n",
       "          const element = document.querySelector('#df-71a9d602-556d-4ab3-9618-c6628552dd1a');\n",
       "          const dataTable =\n",
       "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                     [key], {});\n",
       "          if (!dataTable) return;\n",
       "\n",
       "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "            + ' to learn more about interactive tables.';\n",
       "          element.innerHTML = '';\n",
       "          dataTable['output_type'] = 'display_data';\n",
       "          await google.colab.output.renderOutput(dataTable, element);\n",
       "          const docLink = document.createElement('div');\n",
       "          docLink.innerHTML = docLinkHtml;\n",
       "          element.appendChild(docLink);\n",
       "        }\n",
       "      </script>\n",
       "    </div>\n",
       "  </div>\n",
       "  "
      ],
      "text/plain": [
       "    name   ds  epochs     lr     aug    wd   do   sd  best_val  final_val  \\\n",
       "0  Ti/16  i1k   300.0  0.001  light0  0.03  0.1  0.1  0.702544   0.702232   \n",
       "1  Ti/16  i1k   300.0  0.001  light0  0.03  0.1  0.1  0.702544   0.702232   \n",
       "2  Ti/16  i1k   300.0  0.001  light0  0.03  0.1  0.1  0.702544   0.702232   \n",
       "3  Ti/16  i1k   300.0  0.001  light0  0.03  0.1  0.1  0.702544   0.702232   \n",
       "4  Ti/16  i1k   300.0  0.001  light0  0.03  0.1  0.1  0.702544   0.702232   \n",
       "\n",
       "   ...      adapt_ds adapt_lr  adapt_steps  adapt_resolution  adapt_final_val  \\\n",
       "0  ...  imagenet2012    0.030        20000               384         0.755698   \n",
       "1  ...  imagenet2012    0.010        20000               384         0.754605   \n",
       "2  ...      cifar100    0.030        10000               384         0.836000   \n",
       "3  ...      cifar100    0.010        10000               384         0.835000   \n",
       "4  ...      cifar100    0.003        10000               384         0.800000   \n",
       "\n",
       "   adapt_final_test     params  infer_samples_per_sec  \\\n",
       "0           0.72874  5790000.0                 609.58   \n",
       "1           0.72412  5790000.0                 609.58   \n",
       "2           0.83380  5790000.0                 609.58   \n",
       "3           0.83040  5790000.0                 609.58   \n",
       "4           0.79620  5790000.0                 609.58   \n",
       "\n",
       "                                            filename  \\\n",
       "0  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...   \n",
       "1  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...   \n",
       "2  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...   \n",
       "3  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...   \n",
       "4  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...   \n",
       "\n",
       "                                      adapt_filename  \n",
       "0  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...  \n",
       "1  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...  \n",
       "2  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...  \n",
       "3  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...  \n",
       "4  Ti_16-i1k-300ep-lr_0.001-aug_light0-wd_0.03-do...  \n",
       "\n",
       "[5 rows x 21 columns]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with tf.io.gfile.GFile(\"gs://vit_models/augreg/index.csv\") as f:\n",
    "    df = pd.read_csv(f)\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 820
    },
    "id": "F2Op2iaglEG3",
    "outputId": "55dfcc07-e34e-41ed-cf94-ac986bd69158"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "  <div id=\"df-bcbbb4ac-f81c-4c65-924d-3463793364cd\">\n",
       "    <div class=\"colab-df-container\">\n",
       "      <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>ds</th>\n",
       "      <th>epochs</th>\n",
       "      <th>lr</th>\n",
       "      <th>aug</th>\n",
       "      <th>wd</th>\n",
       "      <th>do</th>\n",
       "      <th>sd</th>\n",
       "      <th>best_val</th>\n",
       "      <th>final_val</th>\n",
       "      <th>...</th>\n",
       "      <th>adapt_ds</th>\n",
       "      <th>adapt_lr</th>\n",
       "      <th>adapt_steps</th>\n",
       "      <th>adapt_resolution</th>\n",
       "      <th>adapt_final_val</th>\n",
       "      <th>adapt_final_test</th>\n",
       "      <th>params</th>\n",
       "      <th>infer_samples_per_sec</th>\n",
       "      <th>filename</th>\n",
       "      <th>adapt_filename</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>32844</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i21k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>none</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.415176</td>\n",
       "      <td>0.414551</td>\n",
       "      <td>...</td>\n",
       "      <td>imagenet2012</td>\n",
       "      <td>0.03</td>\n",
       "      <td>20000</td>\n",
       "      <td>224</td>\n",
       "      <td>0.781299</td>\n",
       "      <td>0.73754</td>\n",
       "      <td>5720000.0</td>\n",
       "      <td>3097.42</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_...</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>32845</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i21k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>none</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.415176</td>\n",
       "      <td>0.414551</td>\n",
       "      <td>...</td>\n",
       "      <td>imagenet2012</td>\n",
       "      <td>0.01</td>\n",
       "      <td>20000</td>\n",
       "      <td>224</td>\n",
       "      <td>0.774586</td>\n",
       "      <td>0.73410</td>\n",
       "      <td>5720000.0</td>\n",
       "      <td>3097.42</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_...</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>32981</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i21k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>none</td>\n",
       "      <td>0.10</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.412598</td>\n",
       "      <td>0.412529</td>\n",
       "      <td>...</td>\n",
       "      <td>imagenet2012</td>\n",
       "      <td>0.01</td>\n",
       "      <td>20000</td>\n",
       "      <td>224</td>\n",
       "      <td>0.772635</td>\n",
       "      <td>0.73368</td>\n",
       "      <td>5720000.0</td>\n",
       "      <td>3097.42</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0...</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10302</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i21k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>light1</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.397148</td>\n",
       "      <td>0.396895</td>\n",
       "      <td>...</td>\n",
       "      <td>imagenet2012</td>\n",
       "      <td>0.03</td>\n",
       "      <td>20000</td>\n",
       "      <td>224</td>\n",
       "      <td>0.767093</td>\n",
       "      <td>0.73106</td>\n",
       "      <td>5720000.0</td>\n",
       "      <td>3097.42</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-d...</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-d...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9826</th>\n",
       "      <td>Ti/16</td>\n",
       "      <td>i21k</td>\n",
       "      <td>300.0</td>\n",
       "      <td>0.001</td>\n",
       "      <td>light0</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.407559</td>\n",
       "      <td>0.407314</td>\n",
       "      <td>...</td>\n",
       "      <td>imagenet2012</td>\n",
       "      <td>0.03</td>\n",
       "      <td>20000</td>\n",
       "      <td>224</td>\n",
       "      <td>0.768810</td>\n",
       "      <td>0.73038</td>\n",
       "      <td>5720000.0</td>\n",
       "      <td>3097.42</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_light0-wd_0.03-d...</td>\n",
       "      <td>Ti_16-i21k-300ep-lr_0.001-aug_light0-wd_0.03-d...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 21 columns</p>\n",
       "</div>\n",
       "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-bcbbb4ac-f81c-4c65-924d-3463793364cd')\"\n",
       "              title=\"Convert this dataframe to an interactive table.\"\n",
       "              style=\"display:none;\">\n",
       "        \n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "       width=\"24px\">\n",
       "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
       "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
       "  </svg>\n",
       "      </button>\n",
       "      \n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      flex-wrap:wrap;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "      <script>\n",
       "        const buttonEl =\n",
       "          document.querySelector('#df-bcbbb4ac-f81c-4c65-924d-3463793364cd button.colab-df-convert');\n",
       "        buttonEl.style.display =\n",
       "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "        async function convertToInteractive(key) {\n",
       "          const element = document.querySelector('#df-bcbbb4ac-f81c-4c65-924d-3463793364cd');\n",
       "          const dataTable =\n",
       "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                     [key], {});\n",
       "          if (!dataTable) return;\n",
       "\n",
       "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "            + ' to learn more about interactive tables.';\n",
       "          element.innerHTML = '';\n",
       "          dataTable['output_type'] = 'display_data';\n",
       "          await google.colab.output.renderOutput(dataTable, element);\n",
       "          const docLink = document.createElement('div');\n",
       "          docLink.innerHTML = docLinkHtml;\n",
       "          element.appendChild(docLink);\n",
       "        }\n",
       "      </script>\n",
       "    </div>\n",
       "  </div>\n",
       "  "
      ],
      "text/plain": [
       "        name    ds  epochs     lr     aug    wd   do   sd  best_val  \\\n",
       "32844  Ti/16  i21k   300.0  0.001    none  0.03  0.0  0.0  0.415176   \n",
       "32845  Ti/16  i21k   300.0  0.001    none  0.03  0.0  0.0  0.415176   \n",
       "32981  Ti/16  i21k   300.0  0.001    none  0.10  0.0  0.0  0.412598   \n",
       "10302  Ti/16  i21k   300.0  0.001  light1  0.03  0.0  0.0  0.397148   \n",
       "9826   Ti/16  i21k   300.0  0.001  light0  0.03  0.0  0.0  0.407559   \n",
       "\n",
       "       final_val  ...      adapt_ds adapt_lr  adapt_steps  adapt_resolution  \\\n",
       "32844   0.414551  ...  imagenet2012     0.03        20000               224   \n",
       "32845   0.414551  ...  imagenet2012     0.01        20000               224   \n",
       "32981   0.412529  ...  imagenet2012     0.01        20000               224   \n",
       "10302   0.396895  ...  imagenet2012     0.03        20000               224   \n",
       "9826    0.407314  ...  imagenet2012     0.03        20000               224   \n",
       "\n",
       "       adapt_final_val  adapt_final_test     params  infer_samples_per_sec  \\\n",
       "32844         0.781299           0.73754  5720000.0                3097.42   \n",
       "32845         0.774586           0.73410  5720000.0                3097.42   \n",
       "32981         0.772635           0.73368  5720000.0                3097.42   \n",
       "10302         0.767093           0.73106  5720000.0                3097.42   \n",
       "9826          0.768810           0.73038  5720000.0                3097.42   \n",
       "\n",
       "                                                filename  \\\n",
       "32844  Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_...   \n",
       "32845  Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_...   \n",
       "32981  Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0...   \n",
       "10302  Ti_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-d...   \n",
       "9826   Ti_16-i21k-300ep-lr_0.001-aug_light0-wd_0.03-d...   \n",
       "\n",
       "                                          adapt_filename  \n",
       "32844  Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_...  \n",
       "32845  Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_...  \n",
       "32981  Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.1-do_0...  \n",
       "10302  Ti_16-i21k-300ep-lr_0.001-aug_light1-wd_0.03-d...  \n",
       "9826   Ti_16-i21k-300ep-lr_0.001-aug_light0-wd_0.03-d...  \n",
       "\n",
       "[5 rows x 21 columns]"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_df = df.query(\n",
    "    f'ds==\"i21k\" & adapt_resolution==224 & adapt_ds==\"imagenet2012\" & name==\"{model_to_convert[0]}\"'\n",
    ").sort_values(\"adapt_final_test\", ascending=False)\n",
    "model_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "v6OecFq6lFSR",
    "outputId": "ed1b6f7f-7e5b-46e9-9aba-1cf17b859056"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224',\n",
       " 0.73754)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_model_i1k_checkpoint = str(model_df.iloc[0][\"adapt_filename\"])\n",
    "model_df.iloc[0][\"adapt_filename\"], model_df.iloc[0][\"adapt_final_test\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ckhktdzblGED",
    "outputId": "a5bb2aa9-60bf-4ab7-f6a3-c14829985f1c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "21.9 MiB - gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz\n"
     ]
    }
   ],
   "source": [
    "filename = best_model_i1k_checkpoint\n",
    "\n",
    "path = f\"gs://vit_models/augreg/{filename}.npz\"\n",
    "\n",
    "print(f\"{tf.io.gfile.stat(path).length / 1024 / 1024:.1f} MiB - {path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "id": "LVpKaY2jlGnT",
    "outputId": "f1e6b1de-594b-4646-a8d0-b882f2051e84"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz'"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "local_path = path.split(\"//\")[-1].split(\"/\")[-1]\n",
    "local_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uf5mypPOnQ9t",
    "outputId": "af04669f-cbb8-462c-89de-4e303141763f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copying gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0--imagenet2012-steps_20k-lr_0.03-res_224.npz...\n",
      "| [1 files][ 21.9 MiB/ 21.9 MiB]                                                \n",
      "Operation completed over 1 objects/21.9 MiB.                                     \n"
     ]
    }
   ],
   "source": [
    "!gsutil cp {path} ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "u49o0iDslHQS"
   },
   "outputs": [],
   "source": [
    "with open(local_path, \"rb\") as f:\n",
    "    params_jax = np.load(f)\n",
    "    params_jax = dict(zip(params_jax.keys(), params_jax.values()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_f6xt9-HnfYx",
    "outputId": "fddf4cc9-d8b4-4149-8cec-ebf2870d3dac"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Transformer/encoder_norm/bias',\n",
      " 'Transformer/encoder_norm/scale',\n",
      " 'Transformer/encoderblock_0/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_0/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_0/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_0/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_0/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_0/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_0/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_0/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_0/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_1/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_1/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_1/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_1/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_1/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_1/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_1/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_1/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_1/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_1/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_1/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_1/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_1/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_1/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_1/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_1/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_10/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_10/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_10/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_10/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_10/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_10/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_10/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_10/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_10/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_10/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_10/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_10/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_10/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_10/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_10/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_10/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_11/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_11/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_11/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_11/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_11/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_11/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_11/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_11/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_11/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_11/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_11/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_11/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_11/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_11/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_11/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_11/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_2/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_2/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_2/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_2/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_2/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_2/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_2/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_2/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_2/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_2/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_2/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_2/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_2/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_2/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_2/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_2/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_3/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_3/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_3/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_3/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_3/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_3/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_3/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_3/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_3/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_3/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_3/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_3/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_3/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_3/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_3/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_3/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_4/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_4/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_4/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_4/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_4/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_4/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_4/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_4/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_4/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_4/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_4/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_4/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_4/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_4/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_4/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_4/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_5/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_5/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_5/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_5/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_5/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_5/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_5/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_5/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_5/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_5/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_5/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_5/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_5/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_5/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_5/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_5/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_6/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_6/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_6/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_6/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_6/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_6/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_6/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_6/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_6/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_6/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_6/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_6/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_6/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_6/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_6/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_6/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_7/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_7/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_7/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_7/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_7/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_7/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_7/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_7/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_7/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_7/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_7/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_7/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_7/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_7/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_7/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_7/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_8/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_8/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_8/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_8/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_8/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_8/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_8/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_8/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_8/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_8/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_8/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_8/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_8/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_8/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_8/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_8/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/encoderblock_9/LayerNorm_0/bias',\n",
      " 'Transformer/encoderblock_9/LayerNorm_0/scale',\n",
      " 'Transformer/encoderblock_9/LayerNorm_2/bias',\n",
      " 'Transformer/encoderblock_9/LayerNorm_2/scale',\n",
      " 'Transformer/encoderblock_9/MlpBlock_3/Dense_0/bias',\n",
      " 'Transformer/encoderblock_9/MlpBlock_3/Dense_0/kernel',\n",
      " 'Transformer/encoderblock_9/MlpBlock_3/Dense_1/bias',\n",
      " 'Transformer/encoderblock_9/MlpBlock_3/Dense_1/kernel',\n",
      " 'Transformer/encoderblock_9/MultiHeadDotProductAttention_1/key/bias',\n",
      " 'Transformer/encoderblock_9/MultiHeadDotProductAttention_1/key/kernel',\n",
      " 'Transformer/encoderblock_9/MultiHeadDotProductAttention_1/out/bias',\n",
      " 'Transformer/encoderblock_9/MultiHeadDotProductAttention_1/out/kernel',\n",
      " 'Transformer/encoderblock_9/MultiHeadDotProductAttention_1/query/bias',\n",
      " 'Transformer/encoderblock_9/MultiHeadDotProductAttention_1/query/kernel',\n",
      " 'Transformer/encoderblock_9/MultiHeadDotProductAttention_1/value/bias',\n",
      " 'Transformer/encoderblock_9/MultiHeadDotProductAttention_1/value/kernel',\n",
      " 'Transformer/posembed_input/pos_embedding',\n",
      " 'cls',\n",
      " 'embedding/bias',\n",
      " 'embedding/kernel',\n",
      " 'head/bias',\n",
      " 'head/kernel']\n"
     ]
    }
   ],
   "source": [
    "from pprint import pformat\n",
    "\n",
    "print(pformat(list(params_jax.keys())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "id": "evio6tbbpFAv"
   },
   "outputs": [],
   "source": [
    "jax_params_to_kcv_params = {\n",
    "    \"Transformer/posembed_input/pos_embedding\": \"patch_embedding/embedding/embeddings\",\n",
    "    \"embedding/bias\": \"patch_embedding_1/dense_26/bias\",\n",
    "    \"embedding/kernel\": \"patch_embedding_1/dense_26/kernel\",\n",
    "    \"cls\": \"patch_embedding_1/class_token\",\n",
    "    \"Transformer/encoderblock_0/LayerNorm_0/scale\": \"transformer_encoder_12/layer_normalization_25/gamma\",\n",
    "    \"Transformer/encoderblock_0/LayerNorm_0/bias\": \"transformer_encoder_12/layer_normalization_25/beta\",\n",
    "    \"Transformer/encoderblock_0/LayerNorm_2/scale\": \"transformer_encoder_12/layer_normalization_26/gamma\",\n",
    "    \"Transformer/encoderblock_0/LayerNorm_2/bias\": \"transformer_encoder_12/layer_normalization_26/beta\",\n",
    "    \"Transformer/encoderblock_0/MultiHeadDotProductAttention_1/query/kernel\": \"transformer_encoder_12/multi_head_attention_12/query/kernel\",\n",
    "    \"Transformer/encoderblock_0/MultiHeadDotProductAttention_1/query/bias\": \"transformer_encoder_12/multi_head_attention_12/query/bias\",\n",
    "    \"Transformer/encoderblock_0/MultiHeadDotProductAttention_1/key/kernel\": \"transformer_encoder_12/multi_head_attention_12/key/kernel\",\n",
    "    \"Transformer/encoderblock_0/MultiHeadDotProductAttention_1/key/bias\": \"transformer_encoder_12/multi_head_attention_12/key/bias\",\n",
    "    \"Transformer/encoderblock_0/MultiHeadDotProductAttention_1/value/kernel\": \"transformer_encoder_12/multi_head_attention_12/value/kernel\",\n",
    "    \"Transformer/encoderblock_0/MultiHeadDotProductAttention_1/value/bias\": \"transformer_encoder_12/multi_head_attention_12/value/bias\",\n",
    "    \"Transformer/encoderblock_0/MultiHeadDotProductAttention_1/out/kernel\": \"transformer_encoder_12/multi_head_attention_12/attention_output/kernel\",\n",
    "    \"Transformer/encoderblock_0/MultiHeadDotProductAttention_1/out/bias\": \"transformer_encoder_12/multi_head_attention_12/attention_output/bias\",\n",
    "    \"Transformer/encoderblock_0/MlpBlock_3/Dense_0/kernel\": \"transformer_encoder_12/dense_27/kernel\",\n",
    "    \"Transformer/encoderblock_0/MlpBlock_3/Dense_0/bias\": \"transformer_encoder_12/dense_27/bias\",\n",
    "    \"Transformer/encoderblock_0/MlpBlock_3/Dense_1/kernel\": \"transformer_encoder_12/dense_28/kernel\",\n",
    "    \"Transformer/encoderblock_0/MlpBlock_3/Dense_1/bias\": \"transformer_encoder_12/dense_28/bias\",\n",
    "    # ... other transformer blocks\n",
    "    \"Transformer/encoder_norm/scale\": \"layer_normalization_49/gamma\",\n",
    "    \"Transformer/encoder_norm/bias\": \"layer_normalization_49/beta\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "CqaW3sqk3xAm",
    "outputId": "eec3924d-2231-4592-bf58-f5d7849f7f57"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<keras.engine.input_layer.InputLayer at 0x7f8d8c9ed6d0>,\n",
       " <keras_cv.layers.vit_layers.PatchingAndEmbedding at 0x7f8cc42f3370>,\n",
       " <keras.layers.regularization.dropout.Dropout at 0x7f8cc43e37c0>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc4306c10>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc431edf0>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc4208d90>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc424ad00>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc42ab2b0>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc420e190>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8d287f3700>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc41bbf40>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc4263e80>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc40fc8b0>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc414fca0>,\n",
       " <keras_cv.layers.transformer_encoder.TransformerEncoder at 0x7f8cc4116a60>,\n",
       " <keras.layers.normalization.layer_normalization.LayerNormalization at 0x7f8cc431e9a0>,\n",
       " <keras.layers.core.lambda_layer.Lambda at 0x7f8cc40bc4c0>,\n",
       " <keras.layers.core.dense.Dense at 0x7f8cc422d760>]"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fjpXbu96KdiT",
    "outputId": "a9ffb226-bb1b-4570-e778-c679368de42f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model_1\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " input_2 (InputLayer)        [(None, 224, 224, 3)]     0         \n",
      "                                                                 \n",
      " patching_and_embedding_1 (P  (None, 197, 192)         185664    \n",
      " atchingAndEmbedding)                                            \n",
      "                                                                 \n",
      " dropout_1 (Dropout)         (None, 197, 192)          0         \n",
      "                                                                 \n",
      " transformer_encoder_12 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_13 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_14 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_15 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_16 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_17 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_18 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_19 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_20 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_21 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_22 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " transformer_encoder_23 (Tra  (None, 197, 192)         444864    \n",
      " nsformerEncoder)                                                \n",
      "                                                                 \n",
      " layer_normalization_49 (Lay  (None, 197, 192)         384       \n",
      " erNormalization)                                                \n",
      "                                                                 \n",
      " lambda_1 (Lambda)           (None, 192)               0         \n",
      "                                                                 \n",
      " dense_49 (Dense)            (None, 1000)              193000    \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 5,717,416\n",
      "Trainable params: 5,717,416\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lrAaY1vzgcPq",
    "outputId": "7ef6f1df-89ff-4d1a-ea79-b5bc37d887ef"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 1, 192)\n",
      "(16, 16, 3, 192)\n",
      "(192,)\n",
      "(1, 197, 192)\n",
      "patching_and_embedding_1/class_token:0 (1, 1, 192)\n",
      "patching_and_embedding_1/conv2d_1/kernel:0 (16, 16, 3, 192)\n",
      "patching_and_embedding_1/conv2d_1/bias:0 (192,)\n",
      "patching_and_embedding_1/embedding/embeddings:0 (197, 192)\n"
     ]
    }
   ],
   "source": [
    "# Check shapes for the class token and embedding layers\n",
    "print(params_jax[\"cls\"].shape)\n",
    "print(params_jax[\"embedding/kernel\"].shape)\n",
    "print(params_jax[\"embedding/bias\"].shape)\n",
    "print(params_jax[\"Transformer/posembed_input/pos_embedding\"].shape)\n",
    "\n",
    "for w in model.layers[1].weights:\n",
    "    print(w.name, w.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "j-MU3hTsuO7X",
    "outputId": "b700a2f2-8803-40f3-8d05-73a37ab129d7"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Variable 'UnreadVariable' shape=(197, 192) dtype=float32, numpy=\n",
       "array([[-0.34854868, -0.02977945, -0.07389098, ...,  0.20677236,\n",
       "        -0.01907007,  0.02419145],\n",
       "       [ 0.35433733,  0.07774402, -0.3533324 , ..., -0.00687748,\n",
       "         0.01606411, -0.05084733],\n",
       "       [ 0.43602812,  0.11496755,  0.33741263, ..., -0.01119205,\n",
       "        -0.02812277, -0.02334888],\n",
       "       ...,\n",
       "       [-0.7694909 , -0.07064533, -0.03692904, ...,  0.07896642,\n",
       "         0.0200146 , -0.04633933],\n",
       "       [-0.5500917 , -0.03244696,  0.06329501, ...,  0.10454354,\n",
       "         0.00393277, -0.03529308],\n",
       "       [-0.4606193 ,  0.03806652, -0.8069602 , ...,  0.13817342,\n",
       "         0.07979445, -0.00859057]], dtype=float32)>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Copy PatchingAndEmbedding layer\n",
    "model.layers[1].weights[0].assign(tf.Variable(params_jax[\"cls\"]))\n",
    "model.layers[1].weights[1].assign(tf.Variable(params_jax[\"embedding/kernel\"]))\n",
    "model.layers[1].weights[2].assign(tf.Variable(params_jax[\"embedding/bias\"]))\n",
    "model.layers[1].weights[3].assign(\n",
    "    tf.Variable(\n",
    "        params_jax[\"Transformer/posembed_input/pos_embedding\"].squeeze()\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "--CvoUtfhIXA",
    "outputId": "90f11a2a-0fa2-4118-c1ea-cd4894c04e30"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(192,)\n",
      "(192,)\n",
      "(192,)\n",
      "(192,)\n",
      "(192, 3, 64)\n",
      "(3, 64)\n",
      "(192, 3, 64)\n",
      "(3, 64)\n",
      "(192, 3, 64)\n",
      "(3, 64)\n",
      "(3, 64, 192)\n",
      "(192,)\n",
      "(192, 768)\n",
      "(768,)\n",
      "(768, 192)\n",
      "(192,)\n",
      "transformer_encoder_13/layer_normalization_27/gamma:0 (192,)\n",
      "transformer_encoder_13/layer_normalization_27/beta:0 (192,)\n",
      "transformer_encoder_13/layer_normalization_28/gamma:0 (192,)\n",
      "transformer_encoder_13/layer_normalization_28/beta:0 (192,)\n",
      "transformer_encoder_13/multi_head_attention_13/query/kernel:0 (192, 3, 64)\n",
      "transformer_encoder_13/multi_head_attention_13/query/bias:0 (3, 64)\n",
      "transformer_encoder_13/multi_head_attention_13/key/kernel:0 (192, 3, 64)\n",
      "transformer_encoder_13/multi_head_attention_13/key/bias:0 (3, 64)\n",
      "transformer_encoder_13/multi_head_attention_13/value/kernel:0 (192, 3, 64)\n",
      "transformer_encoder_13/multi_head_attention_13/value/bias:0 (3, 64)\n",
      "transformer_encoder_13/multi_head_attention_13/attention_output/kernel:0 (3, 64, 192)\n",
      "transformer_encoder_13/multi_head_attention_13/attention_output/bias:0 (192,)\n",
      "transformer_encoder_13/dense_27/kernel:0 (192, 768)\n",
      "transformer_encoder_13/dense_27/bias:0 (768,)\n",
      "transformer_encoder_13/dense_28/kernel:0 (768, 192)\n",
      "transformer_encoder_13/dense_28/bias:0 (192,)\n"
     ]
    }
   ],
   "source": [
    "# Check transformer block shapes between JAX and KCV\n",
    "print(params_jax[\"Transformer/encoderblock_4/LayerNorm_0/scale\"].shape)\n",
    "print(params_jax[\"Transformer/encoderblock_4/LayerNorm_0/bias\"].shape)\n",
    "print(params_jax[\"Transformer/encoderblock_4/LayerNorm_2/scale\"].shape)\n",
    "print(params_jax[f\"Transformer/encoderblock_4/LayerNorm_2/bias\"].shape)\n",
    "print(\n",
    "    params_jax[\n",
    "        f\"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/query/kernel\"\n",
    "    ].shape\n",
    ")\n",
    "print(\n",
    "    params_jax[\n",
    "        f\"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/query/bias\"\n",
    "    ].shape\n",
    ")\n",
    "print(\n",
    "    params_jax[\n",
    "        f\"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/key/kernel\"\n",
    "    ].shape\n",
    ")\n",
    "print(\n",
    "    params_jax[\n",
    "        f\"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/key/bias\"\n",
    "    ].shape\n",
    ")\n",
    "print(\n",
    "    params_jax[\n",
    "        f\"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/value/kernel\"\n",
    "    ].shape\n",
    ")\n",
    "print(\n",
    "    params_jax[\n",
    "        f\"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/value/bias\"\n",
    "    ].shape\n",
    ")\n",
    "print(\n",
    "    params_jax[\n",
    "        f\"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/out/kernel\"\n",
    "    ].shape\n",
    ")\n",
    "print(\n",
    "    params_jax[\n",
    "        f\"Transformer/encoderblock_4/MultiHeadDotProductAttention_1/out/bias\"\n",
    "    ].shape\n",
    ")\n",
    "print(params_jax[f\"Transformer/encoderblock_4/MlpBlock_3/Dense_0/kernel\"].shape)\n",
    "print(params_jax[f\"Transformer/encoderblock_4/MlpBlock_3/Dense_0/bias\"].shape)\n",
    "print(params_jax[f\"Transformer/encoderblock_4/MlpBlock_3/Dense_1/kernel\"].shape)\n",
    "print(params_jax[f\"Transformer/encoderblock_4/MlpBlock_3/Dense_1/bias\"].shape)\n",
    "\n",
    "for w in model.layers[4].weights:\n",
    "    print(w.name, w.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "id": "ypsk93snuPBQ"
   },
   "outputs": [],
   "source": [
    "# Copy Transformer Encoders\n",
    "for i in range(model_to_convert[1][1]):\n",
    "    model.layers[3 + i].weights[0].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[f\"Transformer/encoderblock_{i}/LayerNorm_0/scale\"]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[1].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[f\"Transformer/encoderblock_{i}/LayerNorm_0/bias\"]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[2].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[f\"Transformer/encoderblock_{i}/LayerNorm_2/scale\"]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[3].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[f\"Transformer/encoderblock_{i}/LayerNorm_2/bias\"]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[4].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/kernel\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[5].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/query/bias\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[6].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/kernel\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[7].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/key/bias\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[8].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/kernel\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[9].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/value/bias\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[10].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/kernel\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[11].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MultiHeadDotProductAttention_1/out/bias\"\n",
    "            ].reshape(model.layers[3 + i].weights[11].shape)\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[12].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MlpBlock_3/Dense_0/kernel\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[13].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[f\"Transformer/encoderblock_{i}/MlpBlock_3/Dense_0/bias\"]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[14].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[\n",
    "                f\"Transformer/encoderblock_{i}/MlpBlock_3/Dense_1/kernel\"\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "    model.layers[3 + i].weights[15].assign(\n",
    "        tf.Variable(\n",
    "            params_jax[f\"Transformer/encoderblock_{i}/MlpBlock_3/Dense_1/bias\"]\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KbdPydEZiVAU",
    "outputId": "6a65405f-78da-4892-cc97-52820ccd6cd0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(192,)\n",
      "(192,)\n",
      "layer_normalization_49/gamma:0 (192,)\n",
      "layer_normalization_49/beta:0 (192,)\n"
     ]
    }
   ],
   "source": [
    "print(params_jax[\"Transformer/encoder_norm/scale\"].shape)\n",
    "print(params_jax[\"Transformer/encoder_norm/bias\"].shape)\n",
    "\n",
    "for w in model.layers[15].weights:\n",
    "    print(w.name, w.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "O6Hg5P2Wz5C4",
    "outputId": "7c5e4eb9-549d-4a64-c2e8-98f301f058bc"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Variable 'UnreadVariable' shape=(192,) dtype=float32, numpy=\n",
       "array([ 2.8223321e+00,  8.7457395e-01, -3.7481210e-01,  1.3750441e+00,\n",
       "        4.1197315e-02,  4.8303711e-01, -2.6310286e-01,  3.6670661e-01,\n",
       "       -2.1512895e+00,  2.2048175e+00, -1.3401504e+00,  2.3340700e-02,\n",
       "        8.6943895e-01, -2.1459769e-01,  5.6879807e-01,  9.7927111e-01,\n",
       "        4.7838250e-01, -3.2664791e-01, -1.1942422e-01, -1.3508632e+00,\n",
       "       -1.8508923e-01, -2.8242302e+00,  1.5803963e+00,  4.2268095e+00,\n",
       "       -4.0262440e-01,  6.7515630e-01, -9.2384642e-01, -6.8262619e-01,\n",
       "       -7.9978913e-02,  5.8783874e-02, -1.8007267e+00, -4.0809026e-01,\n",
       "        1.0480539e+00,  2.2342615e+00,  2.4652045e+00,  1.9197918e-01,\n",
       "        1.3735850e+00, -4.0193820e-01,  5.7247961e-01,  6.2033898e-01,\n",
       "       -1.1311350e+00,  7.4566154e+00, -1.1646005e+00,  1.2841954e+00,\n",
       "       -4.0062070e-01,  1.6707023e+00, -1.7686105e+00, -3.9713845e-01,\n",
       "        9.1295403e-01, -1.7427791e+00, -7.2540677e-01, -9.1626835e-01,\n",
       "        3.3038983e+00, -3.0454025e-02, -9.9451184e-02, -2.6315117e-01,\n",
       "       -1.2068810e+00,  3.5536389e+00,  1.5191348e-01,  1.5045251e+00,\n",
       "       -1.2892620e-01,  1.8890557e+00,  5.3895217e-01, -5.0677454e-01,\n",
       "        1.9197935e+00,  2.2222616e-01,  6.1623567e-01, -1.4424083e+00,\n",
       "       -1.1107938e+00,  1.2578135e+00,  5.0049716e-01, -7.5738764e-01,\n",
       "       -6.5021288e-01, -3.6028716e-01,  2.4230197e-01,  3.1851548e-01,\n",
       "        5.8792442e-01, -1.0888535e+00,  2.4443486e+00, -1.9666934e+00,\n",
       "        9.8717093e-01, -6.7094758e-02,  1.6067396e+00,  1.8579993e-01,\n",
       "       -1.0706880e+00, -9.9949531e-02, -1.5306330e+00, -9.4630814e-01,\n",
       "        8.6518264e-01,  9.7798985e-01, -1.5240314e+00,  2.7239501e-01,\n",
       "       -1.7920687e+00, -1.5834644e+00, -1.2013766e+00,  7.7283067e-01,\n",
       "        3.1533870e-01, -1.0921988e+00,  6.8746996e-01, -2.4528224e+00,\n",
       "        5.7313776e-01,  9.5289814e-01, -6.3365275e-01, -1.9187561e+00,\n",
       "        1.7862359e+00, -1.1299679e+00,  3.4474045e-01, -5.3573005e-02,\n",
       "       -3.1702507e-01, -1.1293234e+00, -1.2685273e+00,  2.4756165e-01,\n",
       "       -2.6943071e+00,  6.6708469e-01,  3.2359619e+00, -6.8475269e-02,\n",
       "        1.4569744e+00,  3.5732803e-01,  6.2706912e-01, -2.6309019e-02,\n",
       "        1.5768806e+00, -1.3051279e+00, -5.5007422e-01,  1.2647334e+00,\n",
       "        1.9490814e+00, -1.1320531e+00,  5.1382720e-01,  5.6672585e-01,\n",
       "        1.1814115e+00,  6.4363140e-01,  1.5716019e+00,  3.1377834e-01,\n",
       "        8.0725139e-01,  1.2411931e+00,  3.1249480e+00,  4.3567967e-01,\n",
       "       -8.9678854e-01, -1.9411497e-01, -5.6193280e-01,  1.0215161e+00,\n",
       "       -9.7055250e-01,  8.3245683e-01, -3.3527166e-01, -1.4181898e+00,\n",
       "       -1.6735216e+00,  9.3779010e-01,  3.8624477e-01, -3.7750494e-01,\n",
       "       -6.1774945e-01, -1.0992903e+00, -4.6569797e-01, -3.4994921e-01,\n",
       "        1.2566717e+00, -1.2773196e-01, -1.5977631e+00,  6.9052666e-01,\n",
       "        1.0881550e+00, -2.0216112e+00,  1.7384496e+00, -4.1735049e-02,\n",
       "        1.0723175e+00, -6.9307971e-01,  1.1631954e+00, -6.8310934e-01,\n",
       "        7.7588743e-01,  1.5392259e+00,  7.3759824e-01,  9.7433192e-01,\n",
       "       -1.3154752e+00, -3.8899654e-03,  6.4106512e-01,  2.6138628e-01,\n",
       "       -2.0545945e+00, -9.4495362e-01,  1.0169288e+00, -1.4540728e+00,\n",
       "        5.9237814e-01,  4.5571581e-01, -1.7189237e+00, -9.6421438e-01,\n",
       "        5.1397783e-01, -5.5650976e-02,  9.3600541e-01, -8.6875951e-01,\n",
       "       -1.4851208e+00,  1.3140236e-01, -2.4180408e-01,  9.4073695e-01,\n",
       "        1.0342456e-01, -5.1535720e-01, -1.7350663e-01,  3.4123373e+00],\n",
       "      dtype=float32)>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Copy layer norm before class head\n",
    "model.layers[15].weights[0].assign(\n",
    "    tf.Variable(params_jax[\"Transformer/encoder_norm/scale\"])\n",
    ")\n",
    "model.layers[15].weights[1].assign(\n",
    "    tf.Variable(params_jax[\"Transformer/encoder_norm/bias\"])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jyFglaTAigAC",
    "outputId": "2714cb79-207c-4de9-c7e0-19e804d5d350"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(192, 1000)\n",
      "(1000,)\n",
      "dense_49/kernel:0 (192, 1000)\n",
      "dense_49/bias:0 (1000,)\n"
     ]
    }
   ],
   "source": [
    "print(params_jax[\"head/kernel\"].shape)\n",
    "print(params_jax[\"head/bias\"].shape)\n",
    "\n",
    "for w in model.layers[17].weights:\n",
    "    print(w.name, w.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "s5w-iTku0T2z",
    "outputId": "b670c346-8b7e-47a3-decb-78787caab3cf"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Variable 'UnreadVariable' shape=(1000,) dtype=float32, numpy=\n",
       "array([-3.12163169e-03,  7.99638146e-05, -4.45322460e-03, -5.63935703e-03,\n",
       "       -5.19127958e-03,  1.89294806e-03, -3.09978402e-03,  2.00790493e-03,\n",
       "        1.12810067e-03, -4.83769877e-03, -2.97148316e-03, -6.80502411e-03,\n",
       "       -2.76475097e-03, -5.29246219e-03, -5.48043428e-03, -4.45886562e-03,\n",
       "       -1.57492165e-03, -2.84638233e-03, -9.20321385e-04, -5.74108539e-03,\n",
       "       -5.28989173e-03,  2.09719199e-03, -5.66764874e-03,  1.98722444e-03,\n",
       "       -5.97396633e-03, -2.90608662e-03,  1.00772153e-03, -4.28308034e-03,\n",
       "       -1.85368024e-03,  6.58091565e-04, -5.24879014e-03, -3.40827229e-03,\n",
       "        2.74766865e-03, -1.61942351e-03, -1.62760960e-04, -2.39770347e-03,\n",
       "        9.26783402e-03, -4.49899398e-03, -5.08973782e-04,  1.13704715e-04,\n",
       "       -4.74586571e-03,  8.98223298e-05, -3.58532229e-03, -1.97856897e-03,\n",
       "        3.77449498e-04, -2.17942358e-03, -3.71574331e-03,  1.31800561e-03,\n",
       "       -8.61194357e-03, -1.99970347e-03,  7.55758490e-04,  3.65295238e-03,\n",
       "        3.33294185e-04, -2.57378002e-03,  5.90020034e-04, -4.29749442e-03,\n",
       "       -2.87125120e-03, -6.17771037e-03, -5.40740730e-05,  2.63614440e-03,\n",
       "        7.18098693e-03,  1.60400741e-04,  2.99639278e-03,  2.16771965e-03,\n",
       "       -6.34947355e-05,  2.03361874e-03, -1.40721386e-04, -9.50740185e-04,\n",
       "        2.72718607e-04, -6.03803201e-03, -1.82053493e-03, -2.35399650e-03,\n",
       "       -6.75421581e-03,  3.03951954e-03, -2.61766813e-03, -5.27280383e-03,\n",
       "       -4.57854569e-03, -3.94404493e-03,  2.83436524e-03,  1.51748490e-03,\n",
       "       -1.67738006e-03, -4.85934550e-03,  2.02829742e-05, -2.76851631e-03,\n",
       "       -2.34176451e-03,  1.29400496e-03,  2.11267406e-03,  1.78381219e-03,\n",
       "        2.59598950e-03, -6.46469230e-03, -4.55296366e-03, -2.95093632e-03,\n",
       "       -5.41684404e-03,  1.25572574e-03,  4.43610916e-04, -4.83892672e-03,\n",
       "       -1.06362590e-04, -4.23830003e-03, -3.49641964e-03,  1.70978054e-03,\n",
       "       -2.40980997e-03, -1.77044934e-03, -5.98559808e-03,  7.17883231e-04,\n",
       "       -2.63944094e-04, -4.58216108e-03, -3.40931513e-03, -7.61966454e-03,\n",
       "       -3.63845052e-03, -3.67383403e-03, -4.49879840e-03,  3.25067784e-03,\n",
       "        1.00906892e-02,  6.13356102e-03,  1.57674612e-03, -6.48819283e-03,\n",
       "       -3.51622421e-03, -1.58690149e-03,  3.12452810e-03, -3.64478864e-03,\n",
       "       -2.94588273e-03,  3.32739833e-03,  7.54335313e-04,  1.96225243e-04,\n",
       "        7.86591973e-03, -3.44315940e-03,  6.74191955e-03, -7.68542849e-03,\n",
       "       -9.00785555e-04, -3.27142165e-03, -4.18324256e-03, -3.93421529e-03,\n",
       "       -2.78070988e-03, -4.47021378e-03, -1.94124016e-03, -5.01758279e-03,\n",
       "       -7.19285570e-04, -5.90718491e-03, -4.40639304e-03, -3.50839063e-03,\n",
       "       -4.15522698e-03, -3.65849631e-03, -4.60660737e-03, -6.66346448e-03,\n",
       "       -3.74783133e-03,  4.84544504e-03, -3.17017152e-03, -4.83898213e-03,\n",
       "       -4.76376526e-03, -4.96017141e-03,  4.17507580e-03,  7.99305364e-03,\n",
       "       -5.23286080e-03, -4.50286828e-03, -7.44246878e-04,  9.88742802e-04,\n",
       "       -2.27856101e-03, -2.29081561e-04, -1.42375170e-03,  2.73925922e-04,\n",
       "        2.03565485e-03,  7.69051432e-04, -1.80921913e-03,  4.36893152e-03,\n",
       "       -2.33798390e-04, -2.43295054e-03, -4.03429335e-03, -1.82698201e-03,\n",
       "        1.69122510e-03,  5.86037095e-05,  8.36344028e-04,  3.13251372e-03,\n",
       "        1.03693234e-03,  4.52916138e-04,  1.67019875e-03, -1.22353376e-03,\n",
       "        2.78256671e-03, -1.76786806e-03, -1.30601972e-03,  5.80038934e-04,\n",
       "        8.07047996e-04,  2.14273459e-03,  1.74129987e-03,  1.23082590e-03,\n",
       "        4.79148980e-03,  2.86752521e-03,  3.74320010e-03, -4.09450266e-04,\n",
       "        3.01089557e-03,  5.65809617e-03, -2.58137099e-03,  5.20008290e-03,\n",
       "       -2.16185668e-04,  3.72774806e-03, -3.09997657e-03,  3.27889924e-03,\n",
       "        1.71176100e-03,  4.11220500e-03,  1.11049542e-03,  2.75467685e-03,\n",
       "        1.58426363e-03, -2.07744492e-03,  4.29743854e-03,  1.70426478e-03,\n",
       "       -1.11322186e-03, -4.50864376e-04, -5.05309203e-04, -8.00046546e-04,\n",
       "        3.36495740e-03, -4.76017798e-04, -9.31139337e-04,  2.30265362e-03,\n",
       "        2.60633323e-03, -2.62055406e-03,  8.96915968e-04, -1.26703782e-03,\n",
       "        1.25618698e-03, -7.38864590e-04, -1.16181746e-03,  1.36863277e-03,\n",
       "        5.24641771e-04, -1.13774010e-03,  4.65566118e-04,  4.35540918e-03,\n",
       "        1.36949890e-03,  3.66659905e-03,  2.32451851e-03,  3.20669776e-03,\n",
       "       -3.92755913e-03,  1.06935669e-03, -1.89525040e-03,  2.27358821e-03,\n",
       "        2.92772567e-03, -8.79034516e-04,  1.15636468e-03,  5.62502258e-03,\n",
       "        2.05413997e-03, -1.28134212e-03,  2.00249813e-03, -1.62051409e-03,\n",
       "        6.16867933e-03,  2.74668797e-03,  2.78983754e-03,  1.19094329e-03,\n",
       "        1.91712775e-03,  1.41442125e-03,  1.75950979e-03, -4.53722343e-04,\n",
       "        3.02695960e-04,  2.30059028e-03,  3.50722345e-03,  2.82128900e-03,\n",
       "       -4.84808348e-03,  3.33738443e-03,  1.48802774e-03, -1.63461012e-03,\n",
       "        1.69764098e-03, -1.08216319e-03, -1.59754243e-03,  2.97036604e-04,\n",
       "       -1.72671501e-03,  7.43602868e-05, -1.11791200e-03, -1.94101257e-03,\n",
       "        3.88503144e-03,  3.37074162e-03,  9.33218631e-04,  3.44690285e-03,\n",
       "       -1.37212547e-03, -2.39443779e-03, -3.09062400e-03, -2.10639415e-03,\n",
       "        5.26882941e-04,  2.63194786e-03, -3.26603255e-03, -4.11126437e-03,\n",
       "       -3.87631194e-03, -1.48369989e-03, -5.96996490e-03, -2.61170394e-03,\n",
       "        2.64014816e-03,  2.80521438e-03,  4.31128545e-03, -2.39809765e-03,\n",
       "        2.93605262e-04,  3.83619778e-03, -5.22780349e-04, -4.04319435e-05,\n",
       "       -1.53498491e-03, -5.78478305e-03, -2.04591220e-03, -2.09114654e-03,\n",
       "       -2.16305401e-04, -4.57045995e-03, -2.81028263e-03, -1.45501853e-03,\n",
       "       -3.31376493e-03, -5.65104559e-03, -5.01640548e-04, -4.75169545e-05,\n",
       "       -3.46755073e-03, -1.92387810e-03,  4.18264419e-04,  5.37007488e-03,\n",
       "       -2.10555247e-03, -5.65929571e-03,  2.80797109e-03, -1.95021648e-03,\n",
       "        1.08667809e-04, -4.89568396e-04,  8.49804096e-03,  1.25280255e-03,\n",
       "        2.71669892e-03,  2.10161507e-03,  6.90737693e-03, -1.83544389e-03,\n",
       "       -5.63338958e-03, -2.42078700e-03, -1.41423778e-03, -1.47206453e-03,\n",
       "       -4.47055558e-03, -2.05284404e-03, -2.21793470e-03, -3.87089816e-03,\n",
       "       -3.64430086e-03, -4.75952821e-03, -3.23276827e-03,  9.48843372e-05,\n",
       "       -1.10623281e-04, -4.66065202e-03,  1.76503579e-03,  3.98690114e-03,\n",
       "        1.66877301e-03, -3.84738413e-03,  2.30557425e-03, -3.66852595e-03,\n",
       "       -3.13034584e-03,  3.02909431e-03, -5.41002606e-04, -5.91984252e-03,\n",
       "       -4.92442981e-04,  5.71455527e-03, -1.06137035e-04, -6.97324146e-03,\n",
       "       -5.72499726e-03,  1.07796153e-03, -3.18606198e-03, -6.43902784e-03,\n",
       "       -1.71520514e-03, -1.30873034e-03, -1.36050931e-03, -7.00237323e-03,\n",
       "       -3.15798493e-03,  9.76126350e-04, -1.37792539e-03,  4.55697154e-04,\n",
       "        1.00507801e-02,  1.37135456e-03,  5.27903158e-03,  1.45355484e-03,\n",
       "        9.14939330e-04,  4.12920630e-03,  3.87700717e-03,  4.79707000e-04,\n",
       "       -3.59646999e-03, -5.08697378e-03, -5.45934681e-03, -1.94205437e-03,\n",
       "       -4.40772204e-03, -3.37298214e-03,  4.12713765e-04,  5.85388683e-04,\n",
       "       -9.49628360e-04,  1.99295674e-03, -1.57519287e-04, -3.82821215e-03,\n",
       "       -5.43516316e-03, -5.00654533e-05, -4.48423252e-03,  5.75618644e-04,\n",
       "        6.21759379e-03, -1.61270890e-03, -1.74483622e-03, -2.86163716e-03,\n",
       "       -1.89707451e-03,  2.88602337e-03, -2.09021010e-03, -4.87264153e-03,\n",
       "       -5.36652841e-03,  1.30758423e-03,  4.41172859e-03,  2.49185693e-03,\n",
       "       -5.49624208e-03, -8.06194637e-03,  5.77663013e-04,  4.90133278e-03,\n",
       "       -6.03717053e-03, -4.60737152e-03, -1.89858500e-03,  4.61040210e-04,\n",
       "       -1.66243443e-03,  1.40347343e-03, -2.50215572e-03, -5.54243522e-03,\n",
       "       -3.05146025e-03, -1.36973080e-03, -4.76365071e-03, -5.40275651e-04,\n",
       "       -2.44002906e-03,  6.78489497e-03, -7.39703130e-04, -1.52570638e-03,\n",
       "        1.25025511e-02,  4.14367719e-03, -1.68992183e-03,  4.56893863e-03,\n",
       "       -1.04283530e-03,  3.83881270e-03, -8.71171476e-04,  3.29794898e-03,\n",
       "        5.72145102e-04,  5.68168890e-03, -3.31252604e-03, -1.27264357e-03,\n",
       "       -5.47081511e-03, -2.99424771e-03,  7.44508405e-04,  9.16491449e-03,\n",
       "        8.74707103e-03,  2.03424366e-03,  1.05609058e-03, -2.67508416e-03,\n",
       "        7.71979336e-04,  2.01789732e-03, -1.63753098e-03,  3.98835633e-03,\n",
       "        3.50656832e-04,  1.14561431e-03, -1.15142495e-04, -5.35835628e-04,\n",
       "        6.07113261e-03, -2.05166149e-03,  2.16774340e-03, -5.30683389e-03,\n",
       "       -1.46204780e-03, -7.25094171e-04, -1.40556614e-04,  1.05467346e-02,\n",
       "        1.05009349e-02, -2.25245603e-03, -1.19535963e-03, -4.61603806e-04,\n",
       "        1.44939427e-03, -3.91338579e-03, -2.29449011e-03, -7.02226011e-04,\n",
       "        6.94775488e-03,  5.28433360e-03, -1.28467591e-03, -2.94249039e-04,\n",
       "       -7.44486600e-03,  6.71816431e-03,  3.91850201e-03,  9.80752055e-03,\n",
       "        1.60632702e-03, -4.36776667e-04, -6.24259049e-03, -3.74719989e-03,\n",
       "       -1.19135948e-03,  5.39160892e-03,  1.39754312e-03,  2.32786150e-03,\n",
       "        9.41645529e-04,  2.42740475e-03, -3.26262484e-03,  1.34948222e-03,\n",
       "       -1.35724386e-03,  4.62123047e-04,  2.27372628e-03, -1.05649547e-03,\n",
       "        2.44363537e-03,  1.91295822e-03,  2.45970464e-03, -1.93897262e-03,\n",
       "       -3.35538806e-03, -2.63067125e-03, -3.85016087e-03,  5.68148121e-03,\n",
       "        1.00854710e-02,  3.61315976e-03, -1.56380027e-03,  4.10837587e-03,\n",
       "       -6.78351207e-04, -5.75107941e-03,  8.33634380e-03, -6.06216490e-03,\n",
       "       -8.52113217e-03, -2.46928260e-03, -2.86907214e-03,  4.24410030e-03,\n",
       "       -5.36762690e-03,  2.06253049e-03,  4.65771835e-03, -1.53968390e-03,\n",
       "        9.81930643e-04, -6.66793203e-04,  1.04843071e-02,  6.74432435e-04,\n",
       "       -2.11042585e-03, -3.50448559e-03, -2.57111783e-03, -7.32877525e-04,\n",
       "        1.83022884e-03, -2.00065388e-03,  3.27230082e-03,  7.53291976e-03,\n",
       "        3.09354416e-03, -2.19865073e-03,  7.40691740e-03,  3.27383494e-03,\n",
       "       -7.79334409e-03,  8.52006837e-04,  1.68561784e-03,  7.90746789e-03,\n",
       "        1.59628387e-03,  3.81785707e-04, -6.97317650e-04, -2.02040561e-03,\n",
       "        1.77930895e-04, -2.09211395e-03,  3.75730987e-03, -4.72073822e-04,\n",
       "       -1.69415702e-03, -7.27121299e-03, -3.44159198e-04, -6.10178290e-03,\n",
       "       -6.14984520e-03, -5.85891446e-03,  7.74361019e-04,  3.60932783e-03,\n",
       "       -2.78277835e-03,  6.91215741e-03,  1.16915489e-03,  5.04171476e-04,\n",
       "       -1.95609196e-03,  3.25671677e-03,  8.68304749e-04, -4.62740147e-03,\n",
       "       -6.08578138e-03,  5.02295370e-05, -1.79461529e-03, -8.32979800e-04,\n",
       "        1.58441556e-03, -2.53307819e-03, -3.81275848e-03, -2.44687032e-03,\n",
       "       -2.04430558e-04, -1.55232425e-04,  5.66607434e-03,  4.96462733e-03,\n",
       "       -1.96955446e-03, -3.03122494e-03,  2.86502909e-04, -4.88405814e-03,\n",
       "       -5.80897508e-03,  1.86686884e-04,  3.99583165e-04,  6.18502498e-03,\n",
       "       -1.75048877e-03, -5.00983652e-03,  2.12093676e-03, -2.29248917e-03,\n",
       "        1.61549752e-03, -5.31427050e-03,  1.85987039e-03, -8.25035095e-04,\n",
       "       -4.69916034e-03,  9.76072717e-03,  7.14127207e-04, -1.41469482e-03,\n",
       "        1.08284614e-04, -5.12869284e-03,  1.50497060e-03,  8.78728926e-03,\n",
       "        2.62223766e-03,  7.86311692e-04, -3.63500119e-04,  5.60498144e-03,\n",
       "        1.52141985e-03, -1.42018276e-03, -1.74255169e-03,  2.56176048e-04,\n",
       "       -2.72986013e-03,  7.72156240e-03, -3.93173250e-04, -6.32285024e-04,\n",
       "        6.21773070e-03, -6.11824216e-03, -3.74287856e-03,  1.03500253e-02,\n",
       "        1.61486361e-02,  1.00980292e-03,  4.28333238e-04, -6.28164131e-03,\n",
       "       -5.99983556e-04, -2.93890713e-03,  3.17675294e-03, -2.73337914e-03,\n",
       "        1.74172351e-03, -2.00983300e-03,  9.65832092e-04,  5.47688780e-03,\n",
       "       -8.18866212e-03,  3.56823934e-04, -2.31692917e-03, -8.78005289e-04,\n",
       "        1.00772157e-02, -2.11946361e-04,  9.11396276e-03, -8.80129111e-04,\n",
       "        5.50170988e-03,  5.73248952e-04, -5.34004997e-03,  5.39164525e-04,\n",
       "       -7.54966401e-04, -2.46977247e-03,  4.08835150e-03, -1.98896555e-03,\n",
       "       -2.16142111e-03, -2.98636872e-03, -4.57661255e-04,  3.13019799e-03,\n",
       "        1.94340735e-03,  6.99058454e-03,  1.21065998e-03,  1.39172783e-03,\n",
       "        4.92719654e-03,  4.83877957e-03,  1.62271946e-03, -3.17633909e-04,\n",
       "       -1.01424251e-02,  6.13384508e-03, -1.10292307e-03,  6.73639402e-03,\n",
       "        3.32509796e-03, -5.66059723e-03,  9.48792510e-03, -3.59878666e-03,\n",
       "       -5.94132766e-03, -3.53046623e-03,  3.72506236e-03,  1.55838825e-05,\n",
       "        4.48826142e-03,  1.42325042e-03,  3.08086805e-04,  3.05366493e-03,\n",
       "        1.74211594e-03,  5.62810863e-04, -3.05059832e-03, -2.88241135e-04,\n",
       "        1.87161251e-03, -3.50148114e-03, -3.44939996e-03, -1.05373787e-04,\n",
       "        5.47742005e-03,  2.17020675e-03,  3.56668280e-03,  1.25671376e-03,\n",
       "       -3.84667935e-03, -5.58671216e-03, -9.94620059e-05,  6.66678476e-04,\n",
       "        1.36928400e-03, -5.36611537e-04,  5.26455650e-03, -4.49666334e-03,\n",
       "        3.84758599e-03,  5.44173736e-03,  6.02439791e-03, -1.92842609e-03,\n",
       "        1.65064447e-03,  4.97823406e-04, -1.92961749e-03,  2.95560574e-03,\n",
       "        3.60500673e-03, -5.50317345e-03, -2.97716656e-03, -6.48556137e-03,\n",
       "       -2.01456156e-03, -3.23032867e-03, -5.34666888e-03,  5.45262732e-03,\n",
       "        8.96267407e-03,  1.55172113e-03, -2.04694318e-03,  1.26740313e-03,\n",
       "        4.56992490e-03, -3.60091450e-03, -1.16410549e-03,  2.12472281e-03,\n",
       "        1.02770301e-02,  7.65295932e-04, -2.21458264e-03,  2.78097251e-03,\n",
       "        6.33884571e-04,  1.99714117e-03, -4.52636788e-03,  1.55272719e-03,\n",
       "        4.66679549e-03, -3.72364675e-03,  3.97184119e-03, -2.64048809e-04,\n",
       "       -1.69062999e-03, -1.56632951e-03, -7.77877634e-04, -6.32408401e-03,\n",
       "        5.23838541e-03, -4.24064742e-03, -1.05050474e-03, -4.74008592e-03,\n",
       "       -5.42044058e-04, -5.49049955e-03,  8.51075165e-03, -4.57878271e-03,\n",
       "       -2.82106712e-03,  3.73713858e-03,  7.56579160e-04,  2.45622313e-03,\n",
       "        1.35325473e-02, -5.64372959e-03, -3.67221469e-03,  1.19270356e-02,\n",
       "       -1.19719713e-03,  1.28531391e-02, -7.25171389e-03, -2.73150438e-03,\n",
       "       -2.46329606e-03,  2.46845768e-03,  1.28161898e-02, -1.75325316e-03,\n",
       "        3.46545741e-04, -1.18089840e-04,  2.51080631e-03,  3.59986443e-03,\n",
       "        4.59635584e-03, -3.05025326e-03, -9.10508039e-04,  4.01218334e-04,\n",
       "       -1.16842298e-03,  1.04211019e-02,  1.12514896e-03, -1.06671778e-03,\n",
       "       -4.16937575e-04,  1.54059520e-03,  9.99548938e-04, -2.73163500e-03,\n",
       "        2.24327738e-03,  1.03101856e-03,  8.40578601e-03, -3.41205392e-03,\n",
       "       -3.44345137e-03, -1.35618914e-03,  1.70151389e-03,  1.74899658e-04,\n",
       "        4.80419723e-03,  3.30259837e-03,  1.27451099e-03, -1.73161202e-03,\n",
       "       -6.89298671e-04,  9.55791865e-03,  1.52219727e-03,  1.76438596e-04,\n",
       "        8.76853429e-03,  1.68423820e-03,  6.91810332e-04,  2.77227745e-03,\n",
       "       -8.47049057e-04,  1.25456229e-03,  5.49585372e-03, -3.12717305e-03,\n",
       "       -2.40988471e-03, -3.03346966e-03,  4.67429822e-03,  7.00101256e-03,\n",
       "        5.18984254e-03, -1.82591786e-04,  3.60294012e-03,  5.72476676e-03,\n",
       "       -3.39718140e-03,  2.20567686e-03,  8.83371802e-04,  4.96313768e-03,\n",
       "        6.84178341e-03,  1.08853972e-03, -4.18449054e-03, -9.28529946e-04,\n",
       "        4.15395480e-03,  1.04611483e-03, -6.67769636e-04,  7.24229729e-04,\n",
       "       -2.42743967e-03, -2.61742226e-03, -4.76830266e-03, -2.43613264e-03,\n",
       "        3.54246050e-03,  8.05149134e-03, -1.78565411e-03,  4.44629113e-04,\n",
       "        5.61059266e-03, -3.76659376e-03, -3.03535489e-04, -1.16595707e-03,\n",
       "        3.63611686e-03,  8.75257608e-03, -1.76509237e-03, -9.00687184e-04,\n",
       "       -2.34827818e-03, -1.33259397e-03,  4.86806734e-03, -6.94493356e-04,\n",
       "       -6.32101949e-03, -1.39052782e-03,  7.76145724e-04,  1.08983519e-03,\n",
       "        4.41287365e-03, -3.37803853e-03,  7.97237561e-04,  3.70380096e-03,\n",
       "        1.65338605e-03, -2.64092814e-03,  6.56303624e-03,  7.66313518e-04,\n",
       "       -2.56037316e-03, -1.37079833e-03,  3.20144417e-03, -3.99717479e-04,\n",
       "        8.85554310e-03,  7.60064647e-03,  5.07474970e-03, -1.17071206e-03,\n",
       "        6.10126276e-03, -1.27761997e-03,  6.13473938e-04, -1.08083372e-03,\n",
       "        1.55364978e-03,  7.39243766e-03,  7.04855483e-04, -1.05639093e-03,\n",
       "        1.73761218e-03,  8.39955756e-04,  2.21780315e-03,  9.84709803e-03,\n",
       "        1.27033852e-02,  2.35691507e-04, -4.61897464e-04,  1.76661211e-04,\n",
       "       -2.81910924e-03, -3.28184827e-03,  1.99083472e-03,  5.63934166e-03,\n",
       "       -1.57779257e-03,  4.51659365e-03,  5.80515945e-04, -4.50804364e-03,\n",
       "       -1.79745466e-03,  5.29656885e-04, -4.21042205e-05,  4.38601157e-04,\n",
       "        7.88986217e-03, -3.26596946e-03, -8.87629649e-05, -4.35847975e-03,\n",
       "        5.54971304e-03, -4.89160093e-03, -4.54736501e-03, -6.31081639e-04,\n",
       "        8.20437726e-03, -3.52725596e-03, -2.06540851e-03,  2.81285564e-03,\n",
       "        3.59136844e-03, -5.84498688e-04, -5.27460827e-04,  5.15277684e-03,\n",
       "        3.51102644e-04,  5.09247184e-03, -3.88040650e-03,  1.06626574e-03,\n",
       "       -3.64045915e-03,  1.96496095e-03, -7.66192935e-03, -2.49840406e-04,\n",
       "        1.27221819e-03, -2.63676373e-03, -4.42900974e-03, -8.42618058e-04,\n",
       "        1.38489995e-04, -5.73426951e-04,  6.77291490e-03,  9.19125695e-03,\n",
       "       -5.22898836e-03, -3.00012459e-03,  1.97188952e-03, -4.65750432e-04,\n",
       "        7.42488867e-03, -2.53802398e-04,  4.84076823e-04,  2.21423991e-03,\n",
       "       -5.22186374e-03,  5.28759323e-04,  4.99737170e-03,  5.76268928e-03,\n",
       "        5.35162771e-03, -4.31729294e-03, -3.91714880e-03, -2.77951034e-03,\n",
       "       -2.07758904e-03, -1.34241680e-04,  7.16521172e-04,  1.56212680e-03,\n",
       "       -2.89842626e-03,  8.48609488e-03, -4.08923719e-03,  1.58193125e-03,\n",
       "       -1.56496908e-03, -2.83934223e-03, -6.19054120e-03, -4.49432479e-03,\n",
       "       -2.98070488e-03,  4.65579471e-03, -7.20043608e-04, -1.90585235e-03,\n",
       "       -1.00119819e-03, -6.48335740e-03, -1.30621163e-04, -4.22358559e-03,\n",
       "        4.75490745e-03,  1.81562232e-03, -8.94813274e-04,  8.82010427e-05,\n",
       "       -2.49518873e-03, -3.23481718e-03, -6.21572195e-04, -1.00735319e-03,\n",
       "       -2.87730177e-03, -4.05305531e-03, -2.00531399e-03,  4.09160974e-03,\n",
       "        1.01916993e-03,  5.46160627e-05, -3.40686500e-04, -3.23571241e-03,\n",
       "        1.22432224e-03, -1.78310543e-03,  1.08644483e-03, -1.40198984e-03,\n",
       "        3.29971965e-03, -4.59522847e-03, -3.07088415e-03, -5.87664079e-03,\n",
       "       -2.79736635e-03, -2.66397395e-03, -1.66987674e-03, -5.50022023e-03,\n",
       "       -2.88531510e-03, -2.44558486e-03,  6.97892974e-04, -3.75726772e-03,\n",
       "        2.75574857e-03,  1.64435804e-03, -4.23250109e-04,  6.48514694e-03,\n",
       "        2.71812896e-03,  2.09021801e-03, -4.19980567e-03,  6.97265798e-03,\n",
       "       -2.26788805e-03,  4.63343272e-03,  7.77494861e-04, -2.24448950e-03,\n",
       "        1.54516695e-03, -2.93824426e-03,  2.81573203e-03, -4.54613054e-03,\n",
       "       -2.41043977e-03, -2.63840682e-03, -3.98135884e-03,  9.93194617e-03,\n",
       "        3.49460577e-04, -2.79899337e-03,  2.08845991e-03, -2.34756339e-03,\n",
       "       -3.25305923e-03, -5.42193698e-03, -2.63932953e-03,  2.83103931e-04,\n",
       "        6.01064228e-03, -2.75580864e-03,  1.21093905e-02, -2.53733061e-03],\n",
       "      dtype=float32)>"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Copy haed kernel and bias\n",
    "model.layers[17].weights[0].assign(tf.Variable(params_jax[\"head/kernel\"]))\n",
    "model.layers[17].weights[1].assign(tf.Variable(params_jax[\"head/bias\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "id": "4c-gMql6oc-2"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import numpy as np\n",
    "\n",
    "import PIL\n",
    "import urllib\n",
    "\n",
    "\n",
    "def url_to_array(url):\n",
    "    req = urllib.request.urlopen(url)\n",
    "    arr = np.array(bytearray(req.read()), dtype=np.int8)\n",
    "    arr = cv2.imdecode(arr, -1)\n",
    "    arr = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)\n",
    "    arr = cv2.resize(arr, (224, 224))\n",
    "    return arr\n",
    "\n",
    "\n",
    "def preprocess_image(image, label):\n",
    "    image_resized = tf.image.resize(image, (224, 224))\n",
    "    image_resized = tf.cast(image_resized, tf.float32)\n",
    "    image_resized = (image_resized - 127.5) / 127.5\n",
    "    return image_resized, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "id": "Zx-CnpDl1Dm5"
   },
   "outputs": [],
   "source": [
    "cat = \"https://upload.wikimedia.org/wikipedia/commons/thumb/4/4d/Cat_November_2010-1a.jpg/1200px-Cat_November_2010-1a.jpg\"\n",
    "cat_img = url_to_array(cat)\n",
    "cat_img, _ = preprocess_image(cat_img, None)\n",
    "cat_img = tf.expand_dims(cat_img, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-7S3xee402qP",
    "outputId": "c25874a8-0b56-4bad-b35b-ed95163618a6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2022-12-02 02:18:51--  https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt\n",
      "Resolving storage.googleapis.com (storage.googleapis.com)... 142.250.99.128, 173.194.202.128, 173.194.203.128, ...\n",
      "Connecting to storage.googleapis.com (storage.googleapis.com)|142.250.99.128|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 21675 (21K) [text/plain]\n",
      "Saving to: ‘ilsvrc2012_wordnet_lemmas.txt’\n",
      "\n",
      "ilsvrc2012_wordnet_ 100%[===================>]  21.17K  --.-KB/s    in 0s      \n",
      "\n",
      "2022-12-02 02:18:52 (89.2 MB/s) - ‘ilsvrc2012_wordnet_lemmas.txt’ saved [21675/21675]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt -O ilsvrc2012_wordnet_lemmas.txt\n",
    "\n",
    "with open(\"ilsvrc2012_wordnet_lemmas.txt\", \"r\") as f:\n",
    "    lines = f.readlines()\n",
    "imagenet_int_to_str = [line.rstrip() for line in lines]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-HsjPDvX1Rax",
    "outputId": "ca623a88-15f8-4784-81c9-3bd8108ee063"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - 13s 13s/step\n"
     ]
    }
   ],
   "source": [
    "predictions = model.predict(cat_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_BcVBCrpBAvN",
    "outputId": "e7d1f416-fba6-4853-b1f4-5530291da802"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TopKV2(values=<tf.Tensor: shape=(1, 5), dtype=float32, numpy=\n",
       "array([[0.51381004, 0.33035758, 0.11933842, 0.02861441, 0.00126063]],\n",
       "      dtype=float32)>, indices=<tf.Tensor: shape=(1, 5), dtype=int32, numpy=array([[281, 282, 285, 287, 288]], dtype=int32)>)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "top_5 = tf.math.top_k(predictions, k=5, sorted=False)\n",
    "top_5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "id": "1P-jIeAq1egD",
    "outputId": "1b2c358b-b679-4fbf-c4c2-0d30b01650fe"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'tabby, tabby_cat'"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred = np.argmax(predictions)\n",
    "imagenet_int_to_str[int(pred)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 54
    },
    "id": "91oycPqvDxbR",
    "outputId": "76fcaa87-2925-4a68-dbc2-df5fb110b1b3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1/1 [==============================] - 0s 110ms/step\n"
     ]
    },
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "'golden_retriever'"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dog_url = \"https://hips.hearstapps.com/hmg-prod.s3.amazonaws.com/images/dog-puppy-on-garden-royalty-free-image-1586966191.jpg?crop=1.00xw:0.669xh;0,0.190xh&resize=640:*\"\n",
    "dog_img = url_to_array(dog_url)\n",
    "dog_img, _ = preprocess_image(dog_img, None)\n",
    "dog_img = tf.expand_dims(dog_img, 0)\n",
    "predictions = model.predict(dog_img)\n",
    "pred = np.argmax(predictions)\n",
    "imagenet_int_to_str[int(pred)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "id": "FImkkRO4mdFV"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "id": "7FcoE4lKmdH2"
   },
   "outputs": [],
   "source": [
    "model.compile(\n",
    "    \"adam\",\n",
    "    \"sparse_categorical_crossentropy\",\n",
    "    metrics=[\"accuracy\", keras.metrics.SparseTopKCategoricalAccuracy(5)],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 169,
     "referenced_widgets": [
      "67d9d34aa1d74544b5c0dca2ffb6a5e1",
      "9e40376a6d0a4bc1955be19961fd5a9d",
      "0623879f7ee5477b86c3937847dbeb92",
      "4a93536d899e40dd8d50b5b85c568739",
      "1d66c8c76f114139819c72f801399f3f",
      "08db521e53024651bb28a76e1eeb14bc",
      "222c49d23fa1435da50c9b84b84d57f4",
      "4dcee10d07784e02b37179c87de018b2",
      "efaf16abc036430cb69bcde23d8e4ead",
      "b378dcd372ef46e5ab1a8a7ff9f6441d",
      "e86e67d8f9bd4cb08385bb8ceef245f0",
      "73a973fd35904f279aa26e0df2b342a7",
      "2d81bb1ee5724305b3817d8c151d5526",
      "a2453d88af5d4608abcb40d7927eda3b",
      "4de51dde54a1482891586739c8169c4c",
      "15b8a7b201ef492fb5c945b42c343c01",
      "1a3d7c37503f4547a22cc71b58c12e84",
      "216b3b98ad514af1ad6a1a337bacba66",
      "38e340c53c8d4927a3ac53460eb7345c",
      "830d4ee93daf41e5ac4432cf2b0f65bc",
      "4c3fd7d004de49379daa6c5f9df896f1",
      "63683f3bab43496e8f6b1b904e9deaa5",
      "d323a1e7e9724f068f5ae0dfdc187864",
      "045c83fc9c9947a1a4e8b7b664c0c345",
      "344d084304ba48eaae8101907188aa00",
      "4db4deec686248d4b59f9f0515302b8a",
      "45416111b6fb4115b3000fe20931099f",
      "7947bbec2116428e9ac76302bccb6416",
      "bd72bccb03bf4beb8f928c0ba94883dd",
      "9105f61c122a4f319e10e720ee576c33",
      "1cb29f927d5b4f3e9d6380871ac20dd5",
      "be21c04295ab4856922fe6204acf0274",
      "2ba27d0a983b40a581f42a8db006c018",
      "fdc33137623046fbb7bd1260d57d15a7",
      "68716aba80cd4172aac256cf88bb1bbe",
      "9085b60f6d5844448b0f31ba08bae089",
      "9f6b9b5015a642e1ab53f88b01df7181",
      "236b9d919bb141d5ae6234e9de81fd84",
      "341de77e8e274277b2f58473306dd43f",
      "54aa1bb2cdbe449fb5185e20fc5ad153",
      "7ebc087896c4447b94a8954e3d21ee9c",
      "2a26c33a98d14367a39db44a9600b4e9",
      "916ad258f42d47678e6ae1a579283a59",
      "4cfed962b50441ffaa80354627fbd013",
      "42e7fd74e5424f6ba936d1f8db992f96",
      "528336aa5b204f4ab182c7b12e35c3cc",
      "a8d7d51014ed4d999f7f26999aaf0882",
      "8e588d8f3c8f40f38394d7ba295a0700",
      "e0874ed49fff48868a2b7640ed480c4f",
      "755090c1b6d24829b969cc827e29ec0e",
      "242ca26fbe404449a6c72bf93f22f5da",
      "a558f4e9d50248139931ab5288388cf4",
      "50e3a96fad784e11b7fef267927230da",
      "2f6a123b8be8410aa84bb8e7eeef6071",
      "3727d0fd10d34ea6989f3f535f62acc5",
      "8b272d7b2e3a4b23a2c55197b745a915",
      "84e4d098166f4e5b9e6cc05b52f8ea07",
      "f97b4e321143436dbe8c589715e8faee",
      "1b77b9751d444bae940c8fbf90135a55",
      "cb52127716954b80a279bb6e4aaa3e43",
      "d93b0e9c5b234bf6a2cfc2d746083fd6",
      "6885a2419fb34832bf76d0e52cf6e36b",
      "b9b20b7f7821496eba12a9903a5b759d",
      "05c6202ec51248f6a12379d6642d12f3",
      "7ceb0ec40d6541a6b21ef0b5971b5d66",
      "7020ec9dc40b4dee9f17737e882b67bd"
     ]
    },
    "id": "k111f0PgyHXz",
    "outputId": "80eb4ab3-d904-49aa-fb4e-fc3d20f25e58"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset 1.18 GiB (download: 1.18 GiB, generated: 1.16 GiB, total: 2.34 GiB) to ~/tensorflow_datasets/imagenet_v2/matched-frequency/3.0.0...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67d9d34aa1d74544b5c0dca2ffb6a5e1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dl Completed...: 0 url [00:00, ? url/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "73a973fd35904f279aa26e0df2b342a7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dl Size...: 0 MiB [00:00, ? MiB/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d323a1e7e9724f068f5ae0dfdc187864",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extraction completed...: 0 file [00:00, ? file/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fdc33137623046fbb7bd1260d57d15a7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating splits...:   0%|          | 0/1 [00:00<?, ? splits/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "42e7fd74e5424f6ba936d1f8db992f96",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating test examples...:   0%|          | 0/10000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8b272d7b2e3a4b23a2c55197b745a915",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Shuffling ~/tensorflow_datasets/imagenet_v2/matched-frequency/3.0.0.incomplete39QMVU/imagenet_v2-test.tfrecord…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset imagenet_v2 downloaded and prepared to ~/tensorflow_datasets/imagenet_v2/matched-frequency/3.0.0. Subsequent calls will reuse this data.\n"
     ]
    }
   ],
   "source": [
    "import tensorflow_datasets as tfds\n",
    "\n",
    "(test_set), info = tfds.load(\n",
    "    \"imagenet_v2\", split=[\"test\"], as_supervised=True, with_info=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "id": "1sWLaeSdnUAo"
   },
   "outputs": [],
   "source": [
    "test_set = (\n",
    "    test_set[0]\n",
    "    .shuffle(len(test_set[0]))\n",
    "    .map(preprocess_image)\n",
    "    .batch(32)\n",
    "    .prefetch(tf.data.AUTOTUNE)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mD6OzS57plpo",
    "outputId": "7ce1bec2-9a54-4216-8076-a82620f060d3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[750 231 633 730 465 399 316 242 157 466 537 704 311 974 763 132 281 885\n",
      " 900 372 444 221 727 330 254  87 950 307 563 690 544  80], shape=(32,), dtype=int64)\n"
     ]
    }
   ],
   "source": [
    "for entry, label in test_set.take(1):\n",
    "    print(label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ixhjqYEvrIDt",
    "outputId": "56e20c08-410a-4287-b6eb-4656d8ff7443"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "313/313 [==============================] - 58s 89ms/step - loss: 1.6895 - accuracy: 0.6067 - sparse_top_k_categorical_accuracy: 0.8285\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[1.6895182132720947, 0.6067000031471252, 0.828499972820282]"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.evaluate(test_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "id": "aV11iqJsoQ3N"
   },
   "outputs": [],
   "source": [
    "model.save(f\"{model_to_convert[1][0]}.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3dWrYLIwpAZ6"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.0"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "045c83fc9c9947a1a4e8b7b664c0c345": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7947bbec2116428e9ac76302bccb6416",
      "placeholder": "​",
      "style": "IPY_MODEL_bd72bccb03bf4beb8f928c0ba94883dd",
      "value": "Extraction completed...: 100%"
     }
    },
    "05c6202ec51248f6a12379d6642d12f3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "0623879f7ee5477b86c3937847dbeb92": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_4dcee10d07784e02b37179c87de018b2",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_efaf16abc036430cb69bcde23d8e4ead",
      "value": 1
     }
    },
    "08db521e53024651bb28a76e1eeb14bc": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "15b8a7b201ef492fb5c945b42c343c01": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "1a3d7c37503f4547a22cc71b58c12e84": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "1b77b9751d444bae940c8fbf90135a55": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7ceb0ec40d6541a6b21ef0b5971b5d66",
      "placeholder": "​",
      "style": "IPY_MODEL_7020ec9dc40b4dee9f17737e882b67bd",
      "value": " 9805/10000 [00:04&lt;00:00, 1886.58 examples/s]"
     }
    },
    "1cb29f927d5b4f3e9d6380871ac20dd5": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "1d66c8c76f114139819c72f801399f3f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "216b3b98ad514af1ad6a1a337bacba66": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "222c49d23fa1435da50c9b84b84d57f4": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "236b9d919bb141d5ae6234e9de81fd84": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": "hidden",
      "width": null
     }
    },
    "242ca26fbe404449a6c72bf93f22f5da": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "2a26c33a98d14367a39db44a9600b4e9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "2ba27d0a983b40a581f42a8db006c018": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "2d81bb1ee5724305b3817d8c151d5526": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_1a3d7c37503f4547a22cc71b58c12e84",
      "placeholder": "​",
      "style": "IPY_MODEL_216b3b98ad514af1ad6a1a337bacba66",
      "value": "Dl Size...: 100%"
     }
    },
    "2f6a123b8be8410aa84bb8e7eeef6071": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "341de77e8e274277b2f58473306dd43f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "344d084304ba48eaae8101907188aa00": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_9105f61c122a4f319e10e720ee576c33",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_1cb29f927d5b4f3e9d6380871ac20dd5",
      "value": 1
     }
    },
    "3727d0fd10d34ea6989f3f535f62acc5": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "38e340c53c8d4927a3ac53460eb7345c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": "20px"
     }
    },
    "42e7fd74e5424f6ba936d1f8db992f96": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_528336aa5b204f4ab182c7b12e35c3cc",
       "IPY_MODEL_a8d7d51014ed4d999f7f26999aaf0882",
       "IPY_MODEL_8e588d8f3c8f40f38394d7ba295a0700"
      ],
      "layout": "IPY_MODEL_e0874ed49fff48868a2b7640ed480c4f"
     }
    },
    "45416111b6fb4115b3000fe20931099f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "4a93536d899e40dd8d50b5b85c568739": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_b378dcd372ef46e5ab1a8a7ff9f6441d",
      "placeholder": "​",
      "style": "IPY_MODEL_e86e67d8f9bd4cb08385bb8ceef245f0",
      "value": " 1/1 [00:45&lt;00:00, 37.22s/ url]"
     }
    },
    "4c3fd7d004de49379daa6c5f9df896f1": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "4cfed962b50441ffaa80354627fbd013": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "4db4deec686248d4b59f9f0515302b8a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_be21c04295ab4856922fe6204acf0274",
      "placeholder": "​",
      "style": "IPY_MODEL_2ba27d0a983b40a581f42a8db006c018",
      "value": " 1/1 [00:44&lt;00:00, 44.77s/ file]"
     }
    },
    "4dcee10d07784e02b37179c87de018b2": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": "20px"
     }
    },
    "4de51dde54a1482891586739c8169c4c": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_4c3fd7d004de49379daa6c5f9df896f1",
      "placeholder": "​",
      "style": "IPY_MODEL_63683f3bab43496e8f6b1b904e9deaa5",
      "value": " 1205/1205 [00:44&lt;00:00, 32.52 MiB/s]"
     }
    },
    "50e3a96fad784e11b7fef267927230da": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "528336aa5b204f4ab182c7b12e35c3cc": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_755090c1b6d24829b969cc827e29ec0e",
      "placeholder": "​",
      "style": "IPY_MODEL_242ca26fbe404449a6c72bf93f22f5da",
      "value": "Generating test examples...:  99%"
     }
    },
    "54aa1bb2cdbe449fb5185e20fc5ad153": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "63683f3bab43496e8f6b1b904e9deaa5": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "67d9d34aa1d74544b5c0dca2ffb6a5e1": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_9e40376a6d0a4bc1955be19961fd5a9d",
       "IPY_MODEL_0623879f7ee5477b86c3937847dbeb92",
       "IPY_MODEL_4a93536d899e40dd8d50b5b85c568739"
      ],
      "layout": "IPY_MODEL_1d66c8c76f114139819c72f801399f3f"
     }
    },
    "68716aba80cd4172aac256cf88bb1bbe": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_341de77e8e274277b2f58473306dd43f",
      "placeholder": "​",
      "style": "IPY_MODEL_54aa1bb2cdbe449fb5185e20fc5ad153",
      "value": "Generating splits...: 100%"
     }
    },
    "6885a2419fb34832bf76d0e52cf6e36b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "7020ec9dc40b4dee9f17737e882b67bd": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "73a973fd35904f279aa26e0df2b342a7": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_2d81bb1ee5724305b3817d8c151d5526",
       "IPY_MODEL_a2453d88af5d4608abcb40d7927eda3b",
       "IPY_MODEL_4de51dde54a1482891586739c8169c4c"
      ],
      "layout": "IPY_MODEL_15b8a7b201ef492fb5c945b42c343c01"
     }
    },
    "755090c1b6d24829b969cc827e29ec0e": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "7947bbec2116428e9ac76302bccb6416": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "7ceb0ec40d6541a6b21ef0b5971b5d66": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "7ebc087896c4447b94a8954e3d21ee9c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "830d4ee93daf41e5ac4432cf2b0f65bc": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "84e4d098166f4e5b9e6cc05b52f8ea07": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d93b0e9c5b234bf6a2cfc2d746083fd6",
      "placeholder": "​",
      "style": "IPY_MODEL_6885a2419fb34832bf76d0e52cf6e36b",
      "value": "Shuffling ~/tensorflow_datasets/imagenet_v2/matched-frequency/3.0.0.incomplete39QMVU/imagenet_v2-test.tfrecord*...:  98%"
     }
    },
    "8b272d7b2e3a4b23a2c55197b745a915": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_84e4d098166f4e5b9e6cc05b52f8ea07",
       "IPY_MODEL_f97b4e321143436dbe8c589715e8faee",
       "IPY_MODEL_1b77b9751d444bae940c8fbf90135a55"
      ],
      "layout": "IPY_MODEL_cb52127716954b80a279bb6e4aaa3e43"
     }
    },
    "8e588d8f3c8f40f38394d7ba295a0700": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_2f6a123b8be8410aa84bb8e7eeef6071",
      "placeholder": "​",
      "style": "IPY_MODEL_3727d0fd10d34ea6989f3f535f62acc5",
      "value": " 9930/10000 [00:09&lt;00:00, 1075.68 examples/s]"
     }
    },
    "9085b60f6d5844448b0f31ba08bae089": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_7ebc087896c4447b94a8954e3d21ee9c",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_2a26c33a98d14367a39db44a9600b4e9",
      "value": 1
     }
    },
    "9105f61c122a4f319e10e720ee576c33": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": "20px"
     }
    },
    "916ad258f42d47678e6ae1a579283a59": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "9e40376a6d0a4bc1955be19961fd5a9d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_08db521e53024651bb28a76e1eeb14bc",
      "placeholder": "​",
      "style": "IPY_MODEL_222c49d23fa1435da50c9b84b84d57f4",
      "value": "Dl Completed...: 100%"
     }
    },
    "9f6b9b5015a642e1ab53f88b01df7181": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_916ad258f42d47678e6ae1a579283a59",
      "placeholder": "​",
      "style": "IPY_MODEL_4cfed962b50441ffaa80354627fbd013",
      "value": " 1/1 [00:14&lt;00:00, 14.13s/ splits]"
     }
    },
    "a2453d88af5d4608abcb40d7927eda3b": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_38e340c53c8d4927a3ac53460eb7345c",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_830d4ee93daf41e5ac4432cf2b0f65bc",
      "value": 1
     }
    },
    "a558f4e9d50248139931ab5288388cf4": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "a8d7d51014ed4d999f7f26999aaf0882": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a558f4e9d50248139931ab5288388cf4",
      "max": 10000,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_50e3a96fad784e11b7fef267927230da",
      "value": 10000
     }
    },
    "b378dcd372ef46e5ab1a8a7ff9f6441d": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "b9b20b7f7821496eba12a9903a5b759d": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "bd72bccb03bf4beb8f928c0ba94883dd": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "be21c04295ab4856922fe6204acf0274": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "cb52127716954b80a279bb6e4aaa3e43": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": "hidden",
      "width": null
     }
    },
    "d323a1e7e9724f068f5ae0dfdc187864": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_045c83fc9c9947a1a4e8b7b664c0c345",
       "IPY_MODEL_344d084304ba48eaae8101907188aa00",
       "IPY_MODEL_4db4deec686248d4b59f9f0515302b8a"
      ],
      "layout": "IPY_MODEL_45416111b6fb4115b3000fe20931099f"
     }
    },
    "d93b0e9c5b234bf6a2cfc2d746083fd6": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "e0874ed49fff48868a2b7640ed480c4f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": "hidden",
      "width": null
     }
    },
    "e86e67d8f9bd4cb08385bb8ceef245f0": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "efaf16abc036430cb69bcde23d8e4ead": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "f97b4e321143436dbe8c589715e8faee": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_b9b20b7f7821496eba12a9903a5b759d",
      "max": 10000,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_05c6202ec51248f6a12379d6642d12f3",
      "value": 10000
     }
    },
    "fdc33137623046fbb7bd1260d57d15a7": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_68716aba80cd4172aac256cf88bb1bbe",
       "IPY_MODEL_9085b60f6d5844448b0f31ba08bae089",
       "IPY_MODEL_9f6b9b5015a642e1ab53f88b01df7181"
      ],
      "layout": "IPY_MODEL_236b9d919bb141d5ae6234e9de81fd84"
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
